package com.wuwangfu.func;

import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.async.ResultFuture;
import org.apache.flink.streaming.api.functions.async.RichAsyncFunction;
import org.elasticsearch.action.ActionFuture;
import org.elasticsearch.action.get.GetRequest;
import org.elasticsearch.action.get.GetResponse;
import org.elasticsearch.client.transport.TransportClient;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.transport.TransportAddress;
import org.elasticsearch.transport.client.PreBuiltTransportClient;

import java.net.InetAddress;
import java.util.Collections;
import java.util.concurrent.CompletableFuture;
import java.util.function.Supplier;

/**
 * @Description：自定义ElasticsearchAsyncFunction，实现从ES中查询维度数据
 * @Author：jcshen
 * @Date：2023-07-01
 */
public class AsyncElasticSearchFunc extends RichAsyncFunction<String, Tuple2<String,String>> {
    //es查询客户端
    private transient TransportClient client;

    @Override
    public void open(Configuration parameters) throws Exception {
        //在open方法中设置ES相关参数并创建TransportClient
        Settings settings = Settings.builder().put("cluster.name", "docker-cluster").build();
        client = new PreBuiltTransportClient(settings).addTransportAddress(
                new TransportAddress(InetAddress.getByName("hadoop04"),9300)
//                new TransportAddress(InetAddress.getByName("hadoop04"),9300),
//                new TransportAddress(InetAddress.getByName("hadoop04"),9300)
        );

    }

    @Override
    public void asyncInvoke(String key, ResultFuture<Tuple2<String, String>> resultFuture) throws Exception {
        //使用ES客户端，根据查询条件，到指定的index、type中查询数据
        ActionFuture<GetResponse> future = client.get(new GetRequest("bigdata", "user"));
        CompletableFuture.supplyAsync(new Supplier<String>() {
            @Override
            public String get() {
                try {
                    //获取查询并返回
                    GetResponse response = future.get();
                    return response.getSource().get("info").toString();
                } catch (Exception e) {
                    return null;
                }
            }
        }).thenAccept((String result)->{
            resultFuture.complete(Collections.singleton(Tuple2.of(key,result)));
        });

    }

    @Override
    public void close() throws Exception {
        client.close();
    }
}
